#include <stdlib.h>

#include <iostream>
#include <sstream>
#include <limits>
#include <cmath>
#include <random>

#include <Eigen/Core>
#include <Eigen/Dense>

#include "./include/extractShapes.h"

using namespace ExtractShapes;

/// hyper-parameters
// 随机数引擎
std::default_random_engine g_ranEgn;
// 随机数种子
const static int g_ranSeed = 3;
// 为避免提取点云太小,最大尝试次数和最小r^2
const static int g_maxIter = 10;
const static float g_randRadiusSquare = 0.06;   // 0.25
// KDtree的近邻数量
const static float g_neighbors = 100;   // 30
// 如果点云大小多余指定数目的1.5倍则进行裁剪
const static float g_tooMuchFactor = 1.3;


struct InitConfig {
    InitConfig() {
        g_ranEgn.seed(g_ranSeed);
    }
} initIt;



CloudPtr_t removeOutliers(CloudPtr_t cloud) {
    CloudPtr_t cloud_filtered(new Cloud_t);
    // Create the filtering object
    pcl::StatisticalOutlierRemoval<pcl::PointXYZ> sor;
    sor.setInputCloud (cloud);
    sor.setMeanK (50);
    sor.setStddevMulThresh (1.0);
    sor.filter (*cloud_filtered);
    return cloud_filtered;
}


/// TODO: add normals of point cloud.
std::vector<CloudPtr_t> ExtractShapes::extractPlane(CloudPtr_t cloud, CoefficientPtr_t coefficients) {
    // plane segmentation.
    pcl::SACSegmentation<Point_t> segPlane;
    segPlane.setOptimizeCoefficients(true);
    segPlane.setModelType(pcl::SACMODEL_PLANE);
    segPlane.setMethodType(pcl::SAC_RANSAC);
    segPlane.setMaxIterations(1000);
    segPlane.setDistanceThreshold(0.02);

    // extract plane.
    pcl::ExtractIndices<Point_t> extract;
    pcl::PointIndices::Ptr inliers(new pcl::PointIndices);

    segPlane.setInputCloud(cloud);
    segPlane.segment(*inliers, *coefficients);
    if(inliers->indices.size() == 0) {
        std::cerr << "Could not estimate a planar model for the given dataset." << std::endl;
        return {nullptr, cloud};
    }

    CloudPtr_t cloudTmp(new pcl::PointCloud<Point_t>);
    extract.setInputCloud(cloud);
    extract.setIndices(inliers);
    extract.setNegative(false);
    extract.filter(*cloudTmp);

    CloudPtr_t cloudRest(new pcl::PointCloud<Point_t>);
    extract.setNegative(true);
    extract.filter(*cloudRest);

    return {cloudTmp, cloudRest};
}


NormalPtr_t ExtractShapes::estimateNormal(CloudPtr_t cloud) {
    pcl::NormalEstimation<Point_t, pcl::Normal> ne;
    ne.setInputCloud(cloud);

    pcl::search::KdTree<pcl::PointXYZ>::Ptr 
        tree(new pcl::search::KdTree<pcl::PointXYZ>);
    ne.setSearchMethod(tree);
    NormalPtr_t normals(new Normal_t);
    // ne.setRadiusSearch(0.0300005); 
    ne.setKSearch(g_neighbors);
    ne.compute(*normals);
    return normals;
}


BoundaryPtr_t ExtractShapes::findBoundaries(CloudPtr_t cloud, NormalPtr_t normals) {
    BoundaryPtr_t boundaries(new pcl::PointCloud<pcl::Boundary>);
    pcl::BoundaryEstimation<Point_t, pcl::Normal, pcl::Boundary> est;
    est.setInputCloud(cloud);
    est.setInputNormals(normals);
    pcl::search::KdTree<pcl::PointXYZ>::Ptr tree(new pcl::search::KdTree<pcl::PointXYZ>());
    est.setSearchMethod(tree);
    est.setKSearch(g_neighbors);
    // est.setRadiusSearch(0.0308);
    // est.setRadiusSearch(0.01);
    est.compute(*boundaries);
    return boundaries;
}


void ExtractShapes::writeCloud(CloudPtr_t cloud, const std::string& filename) {
    pcl::PCDWriter writer;
    writer.write<Point_t>(filename, *cloud);
}


void ExtractShapes::show(CloudPtr_t cloud, const std::string& cloudName, const Color_t& color) {
    pcl::visualization::PCLVisualizer viewer ("Example");
    // Define R,G,B colors for the point cloud
    pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZ> colorHandler (cloud, color[0], color[1], color[2]);
    // We add the point cloud to the viewer and pass the color handler
    viewer.addPointCloud (cloud, colorHandler, cloudName.c_str());
    viewer.addCoordinateSystem (1.0, cloudName.c_str(), 0);
    viewer.setBackgroundColor(0.05, 0.05, 0.05, 0); // Setting background to a dark grey
    viewer.setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 2, cloudName.c_str());
    while (!viewer.wasStopped ()) { // Display the visualiser until 'q' key is pressed
        viewer.spinOnce ();
    }
}

//  ------------------------------------------------------------------------
//          class Viewer
//  ------------------------------------------------------------------------

void Viewer::addPointCloud(CloudPtr_t cloud, const std::string& cloudName, const Color_t& color) {
    pcl::visualization::PointCloudColorHandlerCustom<Point_t> 
        colorHandler (cloud, color[0], color[1], color[2]);
    // add the point cloud to the viewer and pass the color handler
    viewer_.addPointCloud (cloud, colorHandler, cloudName.c_str());
    viewer_.setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 2, cloudName.c_str());
}


void Viewer::show() {
    viewer_.addCoordinateSystem (1.0);
    while (!viewer_.wasStopped ()) { // Display the visualiser until 'q' key is pressed
        viewer_.spinOnce ();
    }
}

//  ------------------------------------------------------------------------
//          class RectangleCircleExtractor
//  ------------------------------------------------------------------------

float RectangleCircleExtractor::getprojectedZ(float x, float y) {
    return  - (coefficients->values[0] * x
            + coefficients->values[1] * y 
            + coefficients->values[3]   ) 
            / coefficients->values[2];
}

void RectangleCircleExtractor::extractAndEval(evaluation::Evaluator& eval) {
    // Viewer viewer("point cloud");

    CloudPtr_t cloudToExtract(removeOutliers(cloudOriginal));
    int idx = 1;

    int shapesInPlanes = numOfShapes;
    while(shapesInPlanes > 0) {
        // 提取平面
        auto exteacted = extractPlane(cloudToExtract, coefficients);
        // 提取的平面可能size == 0
        if(exteacted[0] == nullptr || exteacted[0]->size() == 0) break;
        CloudPtr_t cloudPlane = exteacted[0];
        cloudToExtract.swap(exteacted[1]);
        
        int shapesInPlane = shapesInPlanes;
        /// TODO: shapesInPlanes的除数指明了尝试在一个平面中寻找的形状数目
        ///     这个系数应该和场景的复杂度相关，平面越多越复杂
        while(shapesInPlane-- > shapesInPlanes / 2 && cloudPlane->size() != 0) {
            // estimate normals
            NormalPtr_t cloudNormal = estimateNormal(cloudPlane);
            
            // find boundaries.
            BoundaryPtr_t boundaries = findBoundaries(cloudPlane, cloudNormal);

            float randRadiusSquare;
            int selectedPointIdx;
            Eigen::Vector3f selectedPoint;
            int maxIter=g_maxIter;
            do {
                // select a point randomly  
                selectedPointIdx = g_ranEgn() % cloudPlane->size();
                float minDistSquare = std::numeric_limits<float>::max();
                selectedPoint <<    cloudPlane->points[selectedPointIdx].x, 
                                    cloudPlane->points[selectedPointIdx].y,
                                    cloudPlane->points[selectedPointIdx].z;

                // caculate max radius.
                for(size_t i=0; i < cloudPlane->size(); ++i) {
                    int flag = static_cast<int>(boundaries->points[i].boundary_point);
                    if(flag == 1) {
                        float distX = selectedPoint[0] - cloudPlane->points[i].x;
                        float distY = selectedPoint[1] - cloudPlane->points[i].y;
                        float distZ = selectedPoint[2] - cloudPlane->points[i].z;
                        minDistSquare = std::min(distX * distX + distY * distY + distZ* distZ, minDistSquare);
                    }
                }
                randRadiusSquare = minDistSquare; // * (static_cast<float>(g_ranEgn()) / g_ranEgn.max());
            } while(randRadiusSquare < g_randRadiusSquare && --maxIter);

            // extract circle.
            /// TODO: 在之前的计算中可以先进行一遍筛选
            /// 将一些不符合的点提前分到cloudRest中
            CloudPtr_t cloudCircle(new Cloud_t);
            CloudPtr_t cloudRest(new Cloud_t);
            for(size_t i=0; i < cloudPlane->size(); ++i) {
                float distX = selectedPoint[0] - cloudPlane->points[i].x;
                float distY = selectedPoint[1] - cloudPlane->points[i].y;
                float distZ = selectedPoint[2] - cloudPlane->points[i].z;
                if(distX * distX + distY * distY + distZ * distZ < randRadiusSquare)
                    cloudCircle->push_back(cloudPlane->points[i]);
                else
                    cloudRest->push_back(cloudPlane->points[i]);
            } 
            // show(cloudCircle, "xx");
            // show(cloudRest, "yy");

            /// ADDED: resample the point cloud if the points is too much
            ///     discard the point cloud if the pointd is too less.

            if(cloudCircle->size() > minShapeSize) {    // if mix the boundary require.
                CloudPtr_t cloudTarget;
                if(false) {  // circle or rectangle.
                    // if circle is too bigger, shrink it.
                    if(cloudCircle->size() > g_tooMuchFactor * minShapeSize) {
                        CloudPtr_t cloudCircleReal(new Cloud_t);
                        randRadiusSquare *= 1.1 *                                                   // avoid to less points.
                                            static_cast<float>(minShapeSize) / cloudCircle->size(); // required size.
                        for(size_t i=0; i < cloudCircle->size(); ++i) {
                            float distX = selectedPoint[0] - cloudCircle->points[i].x;
                            float distY = selectedPoint[1] - cloudCircle->points[i].y;
                            float distZ = selectedPoint[2] - cloudCircle->points[i].z;
                            if(distX * distX + distY * distY + distZ * distZ < randRadiusSquare)
                                cloudCircleReal->push_back(cloudCircle->points[i]);
                            else
                                cloudRest->push_back(cloudCircle->points[i]);
                        }
                        cloudCircle.swap(cloudCircleReal);
                    }
                    cloudTarget = cloudCircle;
                } 
                else {
                    // select a point to be direction of vector u.
                    size_t randPoint = g_ranEgn() % cloudCircle->size();
                    
                    Eigen::Vector3f u, v, k; 
                    // a vector in plane.
                    u <<    cloudCircle->points[randPoint].x - selectedPoint[0], 
                            cloudCircle->points[randPoint].y - selectedPoint[1],
                            cloudCircle->points[randPoint].z - selectedPoint[2];
                    // normal of plane
                    k << coefficients->values[0], coefficients->values[1], coefficients->values[2];
                    u.normalize();
                    k.normalize();
                    v = u.cross(k); // v.normalize();
                    // select side length randomly.
                    float maxLengthAlongU = randRadiusSquare * (static_cast<float>(g_ranEgn()) / g_ranEgn.max());
                    float maxLengthAlongV = std::sqrt(randRadiusSquare - maxLengthAlongU);
                    maxLengthAlongU = std::sqrt(maxLengthAlongU);

                    // adjust rectangle side length to constrain it in to required ranges.
                    float sizeFactor =  std::sqrt(static_cast<float>(minShapeSize) / cloudCircle->size() *          // required size.
                                        static_cast<float>(randRadiusSquare) / (maxLengthAlongU * maxLengthAlongV));// ratio of area.
                    maxLengthAlongU *= sizeFactor;    // avoid to less points.
                    maxLengthAlongV *= sizeFactor;

                    // do ertract rectangle.
                    CloudPtr_t cloudRectangle(new Cloud_t);
                    for(size_t i=0; i < cloudCircle->size(); ++i) {
                        Eigen::Vector3f thsVec;
                        thsVec <<   cloudCircle->points[i].x - selectedPoint[0], 
                                    cloudCircle->points[i].y - selectedPoint[1], 
                                    cloudCircle->points[i].z - selectedPoint[2];
                        float lengthAlongU = u.dot(thsVec);
                        if(std::abs(lengthAlongU) >= maxLengthAlongU) {
                            cloudRest->push_back(cloudCircle->points[i]);
                            continue;
                        }
                        float lengthAlongV = v.dot(thsVec);
                        if(std::abs(lengthAlongV) >= maxLengthAlongV) {
                            cloudRest->push_back(cloudCircle->points[i]);
                            continue;
                        }
                        cloudRectangle->push_back(cloudCircle->points[i]);
                    }
                    cloudTarget = cloudRectangle;
                }
                
                // save circle or rectangle.
                std::stringstream ss; ss << "cloudTarget" << idx++ << ".pcd";
                // viewer.addPointCloud(cloudTarget, ss.str().c_str(), {255, 255, 0});
                writeCloud(cloudTarget, ss.str().c_str());
                eval.addPointCloud(cloudTarget);
            }
            
            cloudPlane.swap(cloudRest);
        }
        
        // std::stringstream ss("cloudPlane"); ss << idx++;
        // viewer.addPointCloud(cloudPlane, ss.str().c_str(), {255, 0, 255});

        shapesInPlanes = std::min(shapesInPlanes-1, shapesInPlane);
    }
    
    // viewer.show();
}
